import numpy as np
import heapq

# 1. 生成数据
def generate_vectors(num_vectors=10, dimensions=2):
    """ 
    生成向量数据
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据库
    """ 
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=10, dimensions=2)

# 2. 构建近邻图
def construct_knn_graph(vectors, k=3):
    """ 
    构造基于图的 k 近邻图
    :param vectors: 数据点
    :param k: 每个点的近邻数量
    :return: 邻接表表示的近邻图
    """ 
    num_vectors = len(vectors)
    graph = {i: [] for i in range(num_vectors)}
    for i in range(num_vectors):
        distances = []
        for j in range(num_vectors):
            if i != j:
                dist = np.linalg.norm(vectors[i] - vectors[j])
                distances.append((dist, j))
        # 获取 k 个最近邻
        nearest_neighbors = heapq.nsmallest(k, distances)
        graph[i] = [neighbor[1] for neighbor in nearest_neighbors]  # 修正索引错误
    return graph

knn_graph = construct_knn_graph(data_vectors, k=3)

# 3. 基于图的搜索算法
def graph_search(query_vector, vectors, graph, start_node=0):
    """ 
    基于图的近邻搜索
    :param query_vector: 查询向量
    :param vectors: 数据库向量
    :param graph: 近邻图
    :param start_node: 起始节点
    :return: 最接近查询点的向量索引
    """ 
    visited = set()
    current_node = start_node
    min_distance = float('inf')
    nearest_node = current_node
    
    while True:
        visited.add(current_node)
        neighbors = graph[current_node]
        updated = False
        
        for neighbor in neighbors:
            if neighbor not in visited:
                distance = np.linalg.norm(query_vector - vectors[neighbor])
                if distance < min_distance:
                    min_distance = distance
                    nearest_node = neighbor
                    current_node = neighbor
                    updated = True
                    break  # 找到更近的点就移动
        
        if not updated:  # 如果没有找到更近的点，则停止搜索
            break

    return nearest_node, min_distance

# 4. 测试搜索算法
query_vector = np.array([0.5, 0.5])  # 查询向量
nearest_node, distance = graph_search(
    query_vector, data_vectors, knn_graph, start_node=0)

# 5. 输出结果
print("向量数据:")
print(data_vectors)
print("\n近邻图（邻接表表示）:")
for node, neighbors in knn_graph.items():
    print(f"节点 {node}: {neighbors}")

print("\n查询向量:", query_vector)
print(f"最近邻节点索引: {nearest_node}, 距离: {distance:.6f}")